import cv2
import os
import dlib
import numpy as np
from scipy import signal
from scipy.fftpack import rfft, irfft, fftfreq
import pylab as plt
from fncs.ROI import ROI

def FFT_window(data, N, window_size=0):
    # 加窗卷积FFT
    feature_img = np.zeros((N - window_size + 1, window_size))
    for i in range(N):
        j = i - window_size + 1
        if j >= 0:
            feature_img[j] = data[j:i + 1]

    ps = np.zeros((10, window_size))
    for ii in range(10):  # feature_img.shape[0]
        a = rfft(feature_img[ii])
        cf_signal = a.copy()
        Wd = fftfreq(cf_signal.size, d=1)
        cf_signal[(Wd < 0)] = 0
        cf_signal[(Wd > 0.2)] = 0
        ps[ii] = irfft(cf_signal)
    return ps

img_list = os.listdir('../data/TMP/')
img_list.sort(key=lambda x: int(x[:-4]))
B, G, R, C = [], [], [], []

# 模型加载
predictor = dlib.shape_predictor('../models/shape_predictor_68_face_landmarks.dat')

for imgs in img_list:
    # Get ROI
    cheek_L, cheek_R, forehead = ROI(imgs, predictor)
    img = cv2.imread('../data/TMP/%s' % imgs, cv2.IMREAD_COLOR)

    pix_num = cheek_L.shape[0]
    # cheek_L区域
    tmp = np.zeros((pix_num, 3))
    for i in range(pix_num):
        tmp[i] = img[cheek_L[i][0], cheek_L[i][1]]
    # 获取cheek_L区域的RGB通道
    # b, g, r = list(b.flatten()), list(g.flatten()), list(r.flatten())
    B.append(sum(tmp[:, 0]) / pix_num)
    G.append(sum(tmp[:, 1]) / pix_num)
    R.append(sum(tmp[:, 2]) / pix_num)

l, N = 30, np.size(img_list)
fps = 25
H = np.zeros(N)
for n in range(N):   # range(N)
    C.append([R[n], G[n], B[n]])
    m = n - l + 1
    if m >= 0:
        # 计算用于Temp Normalization的mul值
        mul_C = sum(np.array(C[m:n+1])) / l
        # Temp Normalization
        temp_Norm = C[m:n+1] / mul_C
        # Projection
        P_p = np.array([[0, 1, -1], [-2, 1, 1]])
        # P_p = np.array([[3, -2, 0], [1.5, 1, -1.5]])
        S = np.dot(temp_Norm, P_p.T)
        # To compute standard deviation of S1 and S2
        mul_S = sum(S) / l
        std_S = np.sqrt(sum((S - mul_S) ** 2) / l)
        # tuning
        h = S[:, 0] + (std_S[0] / std_S[1]) * S[:, 1]
        mul_h = sum(h) / l
        # overlap-adding
        H[m:n+1] += h - mul_h

plt.subplot(221)
plt.plot(range(N), H)

# Bandpass_filter
W = fftfreq(H.size, d=1)
f_signal = rfft(H)
peaks = signal.find_peaks(f_signal, distance=1)

# search for interval of specific frequency (pulse)
# idx = np.argmin(abs(f_signal[peaks[0]] - 1))   # cheeck_R:1; cheeck_L:1; forehead:1.3
# interval = [W[peaks[0][idx]-1], W[peaks[0][idx]+1]]
interval = [W[peaks[0][2]-1], W[peaks[0][2]+1]]

plt.subplot(222)
plt.plot(W, f_signal)
plt.scatter(W[peaks[0]], f_signal[peaks[0]], marker='.', c='red')
plt.scatter(W[peaks[0][2]], f_signal[peaks[0][2]], marker='^', c='blue')



# If our original signal time was in seconds, this is now in Hz
cut_f_signal = f_signal.copy()
cut_f_signal[(W <= interval[0])] = 0   # [0.05, 0.08, 0.12]
cut_f_signal[(W >= interval[1])] = 0   # [0.08, 0.1, 0.14]
plt.subplot(223)
plt.plot(W, cut_f_signal)

pulse_signal = irfft(cut_f_signal)
peaks = signal.find_peaks(pulse_signal, distance=10)
valley = signal.find_peaks(-pulse_signal, distance=10)

tmp1 = np.array([])
p = 0
q = 1
while q < np.size(peaks[0]):
    tmp1 = np.concatenate((tmp1, [peaks[0][q] - peaks[0][p]]))
    p += 1
    q += 1

tmp2 = np.array([])
p = 0
q = 1
while q < np.size(valley[0]):
    tmp2 = np.concatenate((tmp2, [valley[0][q] - valley[0][p]]))
    p += 1
    q += 1

frequency1 = sum(tmp1) / np.size(tmp1)
frequency2 = sum(tmp2) / np.size(tmp2)
pulse = 60 // (((frequency1 + frequency2) / 2) / fps)
print(pulse)

plt.subplot(224)
plt.plot(range(N), pulse_signal)
plt.scatter(peaks[0], pulse_signal[peaks[0]], marker='v', c='red')
plt.scatter(valley[0], pulse_signal[valley[0]], marker='^', c='red')
plt.show()

# pulse_signal = FFT_window(rPPG_signal, N, window_size=25)
# pulse_signal = pulse_signal.ravel()
# plt.plot(range(pulse_signal.shape[0]), pulse_signal)
# peaks = signal.find_peaks(pulse_signal, distance=15)
#
# print(peaks)
# # print(sum(pulse_signal[peaks[0]][np.where(pulse_signal[peaks[0]] >= 0.002)]) / 10)
# # print(sum(pulse_signal[peaks[0]]) / 234)
# print(pulse_signal[peaks[0]])
#
# plt.scatter(peaks[0], pulse_signal[peaks[0]], marker='^', c='red')
# plt.show()
